#include "MoCapMarkerSensorVisualisation.h"

#include <VirtualRobot/Visualization/CoinVisualization/CoinVisualizationFactory.h>
#include <Inventor/nodes/SoSphere.h>

#include <tuple>

using namespace MMM;

MoCapMarkerSensorVisualisation::MoCapMarkerSensorVisualisation(MoCapMarkerSensorPtr sensor, SoSeparator* sceneSep) :
    SensorVisualisation(sceneSep, new SoSeparator()),
    sensor(sensor)
{
    markerSep = new SoSeparator();
    addChildVisualisationSep(markerSep);
    unlabeledMarkerSep = new SoSeparator();
    addChildVisualisationSep(unlabeledMarkerSep);
}

std::string MoCapMarkerSensorVisualisation::getType() {
    return sensor->getType();
}

int MoCapMarkerSensorVisualisation::getPriority() {
    return sensor->getPriority();
}

SoSeparator* MoCapMarkerSensorVisualisation::createLabel(const std::string &label) {
    SoSeparator* labelSep = new SoSeparator;
    Eigen::Vector3f labelVector;
    labelVector << 0, 30, 0;
    Eigen::Matrix4f localPose = Eigen::Matrix4f::Identity();
    localPose.block(0, 3, 3, 1) += labelVector;
    localPose.block(0, 0, 3, 3) *= 2;
    //labelSep->addChild(marked);
    labelSep->addChild(VirtualRobot::CoinVisualizationFactory::getMatrixTransform(localPose));
    labelSep->addChild(VirtualRobot::CoinVisualizationFactory::CreateBillboardText(label));

    return labelSep;
}

SoSeparator* MoCapMarkerSensorVisualisation::createSphere(const Eigen::Vector3f &pose) {
    SoSeparator* sphereSep = new SoSeparator;
    Eigen::Matrix4f globalPose = Eigen::Matrix4f::Identity();
    globalPose.block(0, 3, 3, 1) = pose;
    SoMaterial* marked = new SoMaterial;
    marked->ambientColor.setValue(0.2f, 0.2f, 1.0f);
    marked->diffuseColor.setValue(0.2f, 0.2f, 1.0f);
    marked->specularColor.setValue(0.5f, 0.5f, 0.5f);
    sphereSep->addChild(marked);
    sphereSep->addChild(VirtualRobot::CoinVisualizationFactory::getMatrixTransform(globalPose));
    SoSphere* sphere = new SoSphere;
    sphere->radius = 10;
    sphereSep->addChild(sphere);

    return sphereSep;
}

void MoCapMarkerSensorVisualisation::setPose(SoSeparator* s, const Eigen::Vector3f &pose) {
    Eigen::Matrix4f globalPose = Eigen::Matrix4f::Identity();
    globalPose.block(0, 3, 3, 1) = pose;
    SoMatrixTransform* mt = (SoMatrixTransform*)s->getChild(1);
    SbMatrix m_(reinterpret_cast<SbMat*>(globalPose.data()));
    mt->matrix.setValue(m_);
}

void MoCapMarkerSensorVisualisation::update(float timestep, float delta) {
    update(sensor->getDerivedMeasurement(timestep, delta));
}

void MoCapMarkerSensorVisualisation::update(float timestep) {
    update(sensor->getDerivedMeasurement(timestep));
}

void MoCapMarkerSensorVisualisation::update(MoCapMarkerSensorMeasurementPtr measurement) {
    if (!measurement) return;

    std::map<std::string, Eigen::Vector3f> labeledMarker = measurement->getLabeledMarker();
    for (const auto &marker : labeledMarker) {
        if (markerSphereSeps.find(marker.first) != markerSphereSeps.end()) {
            // found sphere, then set pose and maybe add child to visualisation
            SoSeparator* sphereSep = std::get<0>(markerSphereSeps[marker.first]);
            bool* sphereSepAdded = std::get<1>(markerSphereSeps[marker.first]);
            if (!(*sphereSepAdded)) markerSep->addChild(sphereSep);
            *sphereSepAdded = true;
            setPose(sphereSep, marker.second);
        } else {
            // not found, create pose
            SoSeparator* sphereSep = createSphere(marker.second);
            sphereSep->addChild(createLabel(marker.first));
            sphereSep->ref();
            markerSep->addChild(sphereSep);
            bool* sphereSepAdded = new bool(true);
            markerSphereSeps[marker.first] = std::make_tuple(sphereSep, sphereSepAdded);
        }
    }
    // remove sphere from visualisation if not in current measurement
    for (const auto &markerSphereSep : markerSphereSeps) {
        if (labeledMarker.find(markerSphereSep.first) == labeledMarker.end()) {
            bool* sphereSepAdded = std::get<1>(markerSphereSep.second);
            if (*sphereSepAdded) markerSep->removeChild(std::get<0>(markerSphereSep.second));
            *sphereSepAdded= false;
        }
    }

    // TODO: Efficency
    if (static_cast<std::size_t>(unlabeledMarkerSep->getNumChildren()) != measurement->getUnlabeledMarker().size()) {
        unlabeledMarkerSep->removeAllChildren();

        for (auto marker : measurement->getUnlabeledMarker()) {
            SoSeparator* sphereSep = createSphere(marker);
            unlabeledMarkerSep->addChild(sphereSep);
        }
    } else {
        size_t pos = 0;
        for (const auto &marker : measurement->getUnlabeledMarker())
        {
            setPose((SoSeparator*) unlabeledMarkerSep->getChild(pos), marker);
            pos++;
        }
    }
}
